#include "task_manager.h"

#include "db_worker_thread.h"

#include <QTimer>

#include <stdexcept>
#include <utility>

namespace base::tasking {

namespace {
constexpr int kDefaultWorkerCount = 4;
}

TaskManager::TaskManager(int workerCount,
                         std::shared_ptr<kernel::ILogger> logger,
                         QObject* parent)
    : QObject(parent)
    , m_workerCount(workerCount > 0 ? workerCount : kDefaultWorkerCount)
    , m_logger(std::move(logger))
{
    m_workers.reserve(m_workerCount);
    for (int i = 0; i < m_workerCount; ++i) {
        auto worker = std::make_unique<DbWorkerThread>(this, i);
        worker->start();
        m_workers.emplace_back(std::move(worker));
    }
}

TaskManager::~TaskManager()
{
    stop();
}

void TaskManager::stop()
{
    if (m_stopping.exchange(true)) {
        return;
    }

    {
        QMutexLocker locker(&m_queueMutex);
        m_taskAvailable.wakeAll();
    }

    for (auto& worker : m_workers) {
        if (worker) {
            worker->stop();
            worker->wait();
        }
    }
    m_workers.clear();
}

std::pair<int, int> TaskManager::snapshotCounts() const
{
    QMutexLocker locker(&m_queueMutex);
    return {m_pendingTaskCount, m_activeTaskCount};
}

TaskHandle::Ptr TaskManager::enqueue(TaskCommand::Ptr command)
{
    if (!command) {
        throw std::invalid_argument("TaskCommand cannot be null");
    }

    if (!command->handle) {
        command->handle = TaskHandle::create(command->taskId, command->batchId);
    }

    if (command->taskId.isEmpty()) {
        command->taskId = command->handle->id();
    }

    command->enqueuedAt = QDateTime::currentMSecsSinceEpoch();

    {
        QMutexLocker locker(&m_queueMutex);
        if (m_stopping.load(std::memory_order_relaxed)) {
            throw std::runtime_error("TaskManager is stopping; no further tasks can be enqueued");
        }
        m_queue.push(command);
        m_pendingIndex[command->taskId] = command;
        ++m_pendingTaskCount;
        m_taskAvailable.wakeOne();
    }

    const auto counts = snapshotCounts();
    emit queueLengthChanged(counts.first, counts.second);

    emit taskEnqueued(command->taskId);

    return command->handle;
}

void TaskManager::cancelTask(const QString& taskId)
{
    if (taskId.isEmpty()) {
        return;
    }

    TaskHandle::Ptr handle;

    {
        QMutexLocker locker(&m_queueMutex);
        m_canceledTaskIds.insert(taskId);

        auto pendingIt = m_pendingIndex.find(taskId);
        if (pendingIt != m_pendingIndex.end()) {
            handle = pendingIt->second->handle;
        }

        auto runningIt = m_runningTasks.find(taskId);
        if (runningIt != m_runningTasks.end()) {
            handle = runningIt->second->handle;
        }
    }

    if (handle) {
        handle->requestCancel();
    }
}

void TaskManager::cancelBatch(const QString& batchId)
{
    if (batchId.isEmpty()) {
        return;
    }

    std::vector<TaskHandle::Ptr> affected;

    {
        QMutexLocker locker(&m_queueMutex);
        m_canceledBatchIds.insert(batchId);

        for (auto& [taskId, command] : m_pendingIndex) {
            if (command && command->batchId == batchId) {
                if (command->handle) {
                    affected.push_back(command->handle);
                }
            }
        }

        for (auto& [taskId, command] : m_runningTasks) {
            if (command && command->batchId == batchId) {
                if (command->handle) {
                    affected.push_back(command->handle);
                }
            }
        }
    }

    for (auto& handle : affected) {
        if (handle) {
            handle->requestCancel();
        }
    }
}

TaskCommand::Ptr TaskManager::takeNextCommand()
{
    QMutexLocker locker(&m_queueMutex);

    while (true) {
        if (m_stopping.load(std::memory_order_relaxed)) {
            return nullptr;
        }

        if (!m_queue.empty()) {
            auto command = m_queue.top();
            m_queue.pop();
            --m_pendingTaskCount;
            if (command) {
                m_pendingIndex.erase(command->taskId);
            }

            if (!command) {
                continue;
            }

            if (command->handle && command->handle->isCanceled()) {
                continue;
            }

            if (m_canceledTaskIds.count(command->taskId) > 0 ||
                (!command->batchId.isEmpty() && m_canceledBatchIds.count(command->batchId) > 0)) {
                if (command->handle) {
                    command->handle->requestCancel();
                }
                continue;
            }

            ++m_activeTaskCount;
            m_runningTasks[command->taskId] = command;
            const auto counts = std::make_pair(m_pendingTaskCount, m_activeTaskCount);
            locker.unlock();
            emit queueLengthChanged(counts.first, counts.second);
            return command;
        }

        m_taskAvailable.wait(&m_queueMutex);
    }
}

void TaskManager::onTaskStarted(const TaskCommand::Ptr& command)
{
    if (command && command->handle) {
        command->handle->notifyRunning();
    }
}

void TaskManager::onTaskFinished(const TaskCommand::Ptr& command, TaskResult result)
{
    if (!command) {
        return;
    }

    bool shouldRetry = !command->handle->isCanceled()
        && !result.success
        && command->attempt < command->retryPolicy.maxRetries;

    {
        QMutexLocker locker(&m_queueMutex);
        m_runningTasks.erase(command->taskId);
        --m_activeTaskCount;
    }

    const auto counts = snapshotCounts();
    emit queueLengthChanged(counts.first, counts.second);

    if (shouldRetry) {
        scheduleRetry(command);
        return;
    }

    if (command->handle->isCanceled()) {
        command->handle->notifyCanceled();
        return;
    }

    if (result.success) {
        command->handle->notifySuccess(result);
        emit taskCompleted(command->taskId);
    } else {
        command->handle->notifyFailure(result);
        if (m_logger) {
            m_logger->error(QString("[TaskManager] Task %1 failed: %2").arg(command->taskId, result.errorMessage));
        }
        emit taskFailed(command->taskId, result.errorMessage);
    }
}

void TaskManager::scheduleRetry(const TaskCommand::Ptr& command)
{
    if (!command) {
        return;
    }

    command->attempt += 1;

    const auto delayMs = static_cast<int>(command->retryPolicy.delay.count());

    if (delayMs <= 0) {
        requeue(command);
        return;
    }

    QTimer::singleShot(delayMs, this, [this, command]() {
        requeue(command);
    });
}

void TaskManager::requeue(const TaskCommand::Ptr& command)
{
    if (!command) {
        return;
    }

    command->enqueuedAt = QDateTime::currentMSecsSinceEpoch();

    {
        QMutexLocker locker(&m_queueMutex);
        if (m_stopping.load(std::memory_order_relaxed)) {
            if (command->handle) {
                command->handle->notifyFailure(TaskResult::failureResult(
                    kernel::ErrorCode::InvalidOperation,
                    "TaskManager stopped before retry could run"));
            }
            return;
        }
        m_queue.push(command);
        m_pendingIndex[command->taskId] = command;
        ++m_pendingTaskCount;
        m_taskAvailable.wakeOne();
    }

    const auto counts = snapshotCounts();
    emit queueLengthChanged(counts.first, counts.second);
}

} // namespace base::tasking
